import os, sys
from typing import Dict

parent_path = os.path.abspath(os.path.join(__file__, *([".."] * 4)))
sys.path.insert(0, parent_path)

import traceback
import numpy as np
import socket, json, cv2, yaml, time
from typing import *
from core.task.modules.processors.base_processor import BaseProcessor
from core.utils.PIL_draw import pic_text
from core.utils.visualize import get_color_map_list
from datetime import datetime, timedelta


class RTMPProcessor(BaseProcessor):
    """直播流检测后处理方案"""

    def __init__(
        self,
        keys,
        RTMPProcessor_cls_map,
    ) -> None:
        self.__in = keys["in"]
        self.__out = keys["out"]
        self.__detect = keys["detect"]

        self.__cls_map = {}
        self.__color_map = {}
        index = 0

        # 匹配识别类型,获取类别总数
        for key in self.__detect:
            if key in RTMPProcessor_cls_map:
                if isinstance(RTMPProcessor_cls_map[key], dict):
                    self.__cls_map[key] = RTMPProcessor_cls_map[key]
                elif isinstance(RTMPProcessor_cls_map[key], str):
                    with open(RTMPProcessor_cls_map[key], "r") as f:
                        self.__cls_map[key] = {
                            key_: value_
                            for key_, value_ in enumerate(
                                yaml.safe_load(f)["label_list"]
                            )
                        }
                self.__color_map[key] = index
                index += len(self.__cls_map[key])

        # 初始化颜色列表,确保所有模型的所有类型都是一个单独的颜色
        self.__color_list = get_color_map_list(index)

    def init_check(self):
        assert not self.__in is None
        assert not self.__out is None
        assert not self.__detect is None
        assert not self.__cls_map is None
        assert not self.__color_map is None
        super().init_check()

    def __call__(self, data: Dict) -> Dict:
        # 获取图片数据
        detect_images = data[self.__in][0].copy()
        draw_thickness = min(detect_images.shape[:2]) // 320

        # 是否有检测到目标
        have_target = False

        # 解析数据
        for key in self.__detect:
            if key in data:
                for idx, bboxes in enumerate(data[key]["boxes"][0]):
                    cls_id, score, x1, y1, x2, y2 = bboxes
                    cls_id = int(cls_id)
                    # 过滤检测结果
                    if cls_id not in self.__cls_map[key]:
                        continue

                    # 绘制识别框
                    detect_images = cv2.rectangle(
                        detect_images,
                        (int(x1), int(y1)),
                        (int(x2), int(y2)),
                        color=self.__color_list[cls_id + self.__color_map[key]],
                        thickness=draw_thickness,
                    )
                    # 绘制文字
                    cv2.putText(
                        detect_images,
                        self.__cls_map[key][cls_id],
                        (int(x1), int(y1)),
                        cv2.FONT_HERSHEY_COMPLEX_SMALL,
                        1,
                        color=self.__color_list[cls_id + self.__color_map[key]],
                    )

        # 保存数据
        data[self.__out] = detect_images

        return data

    def close(self):
        self.status = 0


if __name__ == "__main__":
    # 取对应的内部配置
    cfg = {
        "keys": {
            "in": "input_data",
            "out": "visualize",
            "detect": ["bird_detect", "cannon_detect", "dolphin_detect"],
        }
    }

    rtmp_processor = RTMPProcessor(**cfg)
